package com.bruce.ai.alibaba.workflow.config;

import java.util.HashMap;
import java.util.Map;

import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.agent.ReactAgent;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.bruce.ai.alibaba.workflow.node.AgentInvokeNode;
import com.bruce.ai.alibaba.workflow.node.PreprocessorNode;
import com.bruce.ai.alibaba.workflow.node.ValidatorNode;

import static com.alibaba.cloud.ai.graph.action.AsyncEdgeAction.edge_async;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;


/**
 * 混合使用 Agent Node 和 普通Node
 */
@Configuration
public class HybridWorkflowConfig {

    @Bean
    public StateGraph hybridWorkflow(@Qualifier("qaAgent") ReactAgent qaAgent) throws Exception{
        //定义策略路线
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> strategies = new HashMap<>();
            strategies.put("input", new ReplaceStrategy());
            strategies.put("cleaned_input", new ReplaceStrategy());
            strategies.put("qa_result", new ReplaceStrategy());
            strategies.put("is_valid", new ReplaceStrategy());
            return strategies;
        };
        
        // 构建 StateGraph
        StateGraph graph = new StateGraph(keyStrategyFactory);

        // 添加普通节点
        graph.addNode("preprocess", node_async(new PreprocessorNode()));
        graph.addNode("validate", node_async(new ValidatorNode()));

        // 替换为手动调用节点，避免框架内对 messages 的类型假设
        graph.addNode("qa_agent",  node_async(new AgentInvokeNode(qaAgent, "qa_result")));

        // 定义流程
        graph.addEdge(StateGraph.START, "preprocess");
        graph.addEdge("preprocess", "qa_agent");
        graph.addEdge("qa_agent", "validate");

        // 添加条件边  
        graph.addConditionalEdges("validate"
        , edge_async(state -> (Boolean)state.value("is_valid").orElse(false)? "end" : "qa_agent")
        , Map.of("end", StateGraph.END, "qa_agent", "qa_agent"));
        return graph;
    }

}
